#include "./Kalman_Filter_Coor.hpp"

#include <float.h>

#define LOG_TAG "KF coor"    /* 模块TAG */
#define LOG_LVL LOG_LVL_INFO /* 静态过滤级别 */
#include <ulog.h>            /* 必须放在宏定义下面 */

static union Kalman_Filter_Coor_Namespace::Sys_State_Union sys_state_data, sys_state_forecast_data, measure_data; /* 用于储存状态量，预测状态量和测量值 */
static arm_matrix_instance_f32 sys_state_matrix, sys_state_forecast_matrix, measure_matrix;                       /* 状态量矩阵，预测状态量矩阵和测量矩阵 */

static arm_matrix_instance_f32 input_matrix; /* 控制量 */
static float input_matrix_data[3];

static float covariance_data[9], covariance_forecast_data[9];                 /* 用于存储协方差和预测协方差 */
static arm_matrix_instance_f32 covariance_matrix, covariance_forecast_matrix; /* 状态量矩阵和预测状态量矩阵 */

static float W_data[9], R_data[9], H_data[9], Q_DM_data[9], Q_TC_data[9], Q_data[9], V_data[9];            /* 存储噪声 */
static arm_matrix_instance_f32 W_matrix, R_matrix, H_matrix, Q_DM_matrix, Q_TC_matrix, Q_matrix, V_matrix; /* 控制噪声和测量噪声 */

static float K_data[9], I_data[9];                 /* 用于存储卡尔曼增益 */
static arm_matrix_instance_f32 K_matrix, I_matrix; /* 卡尔曼增益矩阵和单位矩阵 */

static float A_data[9], B_data[9];                 /* 状态转换矩阵A，控制转换矩阵B */
static arm_matrix_instance_f32 A_matrix, B_matrix; /* 状态转换矩阵，控制转换矩阵 */

static float Z_temp1_data[9], Z_temp2_data[9], Z_temp3_data[9]; /* 用于h()函数的计算 */
static arm_matrix_instance_f32 Z_temp1_matrix, Z_temp2_matrix, Z_temp3_matrix;

/* 初始化矩阵 */
void Kalman_Filter_Coor_Namespace::Init(float dm_x_noise, float dm_y_noise, float dm_angle_noise, /* 输入量为二维码的测量噪声 */
                                        float ct_x_noise, float ct_angle_noise                    /* 输入量为色带的测量噪声 */
)
{
    arm_mat_init_f32(&sys_state_matrix, 3, 1, sys_state_data.data);
    arm_mat_init_f32(&sys_state_forecast_matrix, 3, 1, sys_state_forecast_data.data);
    arm_mat_init_f32(&measure_matrix, 3, 1, measure_data.data);
    arm_mat_init_f32(&input_matrix, 3, 1, input_matrix_data);
    arm_mat_init_f32(&covariance_matrix, 3, 3, covariance_data);
    arm_mat_init_f32(&covariance_forecast_matrix, 3, 3, covariance_forecast_data);

    arm_mat_init_f32(&W_matrix, 3, 3, W_data);
    arm_mat_init_f32(&R_matrix, 3, 3, R_data);
    arm_mat_init_f32(&H_matrix, 3, 3, H_data);
    arm_mat_init_f32(&Q_DM_matrix, 3, 3, Q_DM_data);
    arm_mat_init_f32(&Q_TC_matrix, 3, 3, Q_TC_data);
    arm_mat_init_f32(&Q_matrix, 3, 3, Q_data);
    arm_mat_init_f32(&V_matrix, 3, 3, V_data);

    arm_mat_init_f32(&K_matrix, 3, 3, K_data);
    arm_mat_init_f32(&I_matrix, 3, 3, I_data);
    arm_mat_init_f32(&A_matrix, 3, 3, A_data);
    arm_mat_init_f32(&B_matrix, 3, 3, B_data);

    arm_mat_init_f32(&Z_temp1_matrix, 3, 3, Z_temp1_data);
    arm_mat_init_f32(&Z_temp2_matrix, 3, 3, Z_temp2_data);
    arm_mat_init_f32(&Z_temp3_matrix, 3, 3, Z_temp3_data);

    int i = 0;
    /* 状态量初始值为0 */
    for (i = 0; i < 3; i++)
    {
        sys_state_data.data[i] = 0.0f;
        sys_state_forecast_data.data[i] = 0.0f;
        measure_data.data[i] = 0.0f;
    }

    for (i = 0; i < 9; i++)
    {
        W_data[i] = 0.0f;
        H_data[i] = 0.0f;
        Q_DM_data[i] = 0.0f;
        Q_TC_data[i] = 0.0f;
        V_data[i] = 0.0f;
    }

    for (i = 0; i < 9; i++)
    {
        K_data[i] = 0.0f;
        I_data[i] = 0.0f;
        B_data[i] = 0.0f;
        A_data[i] = 0.0f;
    }

    for (i = 0; i < 9; i++)
    {
        Z_temp1_data[i] = 0.0f;
        Z_temp2_data[i] = 0.0f;
        Z_temp3_data[i] = 0.0f;
    }

    I_data[0] = I_data[4] = I_data[8] = 1.0f;
    A_data[0] = A_data[4] = A_data[8] = 1.0f;
    /* 上电初始时，全局坐标不确定，因此设置协方差为极大 */
    covariance_data[0] = covariance_data[4] = covariance_data[8] = 1.0e+10f;

    /* 色带无法测量纵向坐标，因此方差为正无穷 */
    Q_TC_data[0] = ct_x_noise;
    Q_TC_data[4] = FLT_MAX;
    Q_TC_data[8] = ct_angle_noise;

    Q_DM_data[0] = dm_x_noise;
    Q_DM_data[4] = dm_y_noise;
    Q_DM_data[8] = dm_angle_noise;

    H_data[8] = 1.0f;
    V_data[0] = V_data[4] = V_data[8] = 1.0f;

    Z_temp1_data[4] = Z_temp1_data[8] = 1.0f;
    Z_temp2_data[8] = 1.0f;
}

/* 使用二维码对坐标滤波*/
void Kalman_Filter_Coor_Namespace::Kalman_Filter_Coor(const union Kalman_Filter_Velocity_Namespace::Sys_State_Union &input, /* 输入控制量 */
                                                      const arm_matrix_instance_f32 &R,                                     /* 控制噪声 */
                                                      const union Sys_State_Union &measure,                                 /* 二维码测量值 */
                                                      const float time_s                                                    /* 间隔时间 */
)
{
    Kalman_Filter_Coor_Namespace::Update_Coor_No_Measurement(input, R, time_s, false);          /* 预测状态量和协方差 */
    Kalman_Filter_Function::Cal_Kalman_Gain(covariance_forecast_matrix, Q_DM_matrix, K_matrix); /* 计算卡尔曼增益 */
    for (int i = 0; i < 3; i++)
    {
        /* 更新测量值 */
        measure_data.data[i] = measure.data[i];
    }
    Kalman_Filter_Function::Update_State_Variable(sys_state_forecast_matrix, K_matrix, measure_matrix, sys_state_matrix); /* 更新状态矩阵 */
    Kalman_Filter_Function::Update_Covariance(covariance_forecast_matrix, I_matrix, K_matrix, covariance_matrix);         /* 更新协方差矩阵 */
    LOG_D("w:%f,theta:%f by dm.", input.w, sys_state_data.theta);
}

/* 使用色带对坐标滤波*/
void Kalman_Filter_Coor_Namespace::Kalman_Filter_Coor(const union Kalman_Filter_Velocity_Namespace::Sys_State_Union &input, /* 输入控制量 */
                                                      const arm_matrix_instance_f32 &R,                                     /* 控制噪声 */
                                                      const union Sys_State_Union &measure,                                 /* 色带测量值 */
                                                      const union Sys_State_Union &coor_base,                               /* 色带基坐标 */
                                                      const float time_s                                                    /* 间隔时间 */
)
{
    Kalman_Filter_Coor_Namespace::Update_Coor_No_Measurement(input, R, time_s, false); /* 预测状态量和协方差 */

    /* 计算正弦余弦 */
    float sin_theta_base, cos_theta_base;
    float sin_theta_z, cos_theta_z;
    sin_theta_base = arm_sin_f32(coor_base.theta);
    cos_theta_base = arm_cos_f32(coor_base.theta);
    sin_theta_z = arm_sin_f32(measure.theta);
    cos_theta_z = arm_cos_f32(measure.theta);
    /* 更新H矩阵 */
    H_data[0] = cos_theta_base / cos_theta_z;
    H_data[1] = sin_theta_base / cos_theta_z;
    H_data[3] = -sin_theta_base;
    H_data[4] = cos_theta_base;
    /* 更新V矩阵 */
    {
        float temp1, temp2, cos_temp;
        cos_temp = cos_theta_z * cos_theta_z;
        temp1 = cos_theta_base * sin_theta_z * (sys_state_forecast_data.x - coor_base.x) / cos_temp;
        temp2 = sin_theta_base * sin_theta_z * (sys_state_forecast_data.y - coor_base.y) / cos_temp;
        V_data[2] = temp1 + temp2;
    }

    Kalman_Filter_Function::Cal_A_mul_B_A_Trans(V_matrix, Q_TC_matrix, Q_matrix);                      /* 计算转换后的测量噪声矩阵 */
    Kalman_Filter_Function::Cal_Kalman_Gain(covariance_forecast_matrix, H_matrix, Q_matrix, K_matrix); /* 计算卡尔曼增益 */

    arm_matrix_instance_f32 sys_temp1_matrix, sys_temp2_matrix;
    union Kalman_Filter_Coor_Namespace::Sys_State_Union sys_temp1_data, sys_temp2_data;
    arm_mat_init_f32(&sys_temp1_matrix, 3, 1, sys_temp1_data.data);
    arm_mat_init_f32(&sys_temp2_matrix, 3, 1, sys_temp2_data.data);
    for (int i = 0; i < 3; i++)
    {
        sys_temp1_data.data[i] = sys_state_forecast_data.data[i] - coor_base.data[i];
    }

    Z_temp1_data[0] = 1.0f / cos_theta_z;
    Z_temp2_data[0] = cos_theta_base;
    Z_temp2_data[1] = sin_theta_base;
    Z_temp2_data[3] = -sin_theta_base;
    Z_temp2_data[4] = cos_theta_base;
    arm_mat_mult_f32(&Z_temp1_matrix, &Z_temp2_matrix, &Z_temp3_matrix);
    arm_mat_mult_f32(&Z_temp3_matrix, &sys_temp1_matrix, &sys_temp2_matrix);

    /* 更新测量值 */
    for (int i = 0; i < 3; i++)
    {
        measure_data.data[i] = measure.data[i];
    }

    {
        arm_matrix_instance_f32 delta_z_matrix;
        union Kalman_Filter_Coor_Namespace::Sys_State_Union delta_z_data;
        arm_mat_init_f32(&delta_z_matrix, 3, 1, delta_z_data.data);
        for (int i = 0; i < 3; i++)
        {
            delta_z_data.data[i] = measure_data.data[i] - sys_temp2_data.data[i];
        }

        /* 分配并计算K*delta_z_matrix */
        float K_mul_delta_z_matrix_data[3];
        arm_matrix_instance_f32 K_mul_delta_z_matrix_matrix;
        arm_mat_init_f32(&K_mul_delta_z_matrix_matrix, 3, 1, K_mul_delta_z_matrix_data);
        arm_mat_mult_f32(&K_matrix, &delta_z_matrix, &K_mul_delta_z_matrix_matrix);

        /* 更新状态量 */
        arm_mat_add_f32(&sys_state_forecast_matrix, &K_mul_delta_z_matrix_matrix, &sys_state_matrix);
    }

    Kalman_Filter_Function::Update_Covariance(covariance_forecast_matrix, I_matrix, K_matrix, H_matrix, covariance_matrix); /* 更新协方差矩阵 */
    LOG_D("w:%f,theta:%f by tc.", input.w, sys_state_data.theta);
}

/* 无外部测量时更新坐标 */
void Kalman_Filter_Coor_Namespace::Update_Coor_No_Measurement(const union Kalman_Filter_Velocity_Namespace::Sys_State_Union &input, /* 输入控制量 */
                                                              const arm_matrix_instance_f32 &R,                                     /* 控制噪声 */
                                                              const float time_s,                                                   /* 间隔时间 */
                                                              const bool no_measurement_flag)
{
    float cos_theta, sin_theta;

    /* 计算正弦余弦 */
    sin_theta = arm_sin_f32(sys_state_data.theta);
    cos_theta = arm_cos_f32(sys_state_data.theta);

    const float cos_mul_time = cos_theta * time_s;
    const float sin_mul_time = sin_theta * time_s;

    B_data[0] = cos_mul_time;
    B_data[1] = -sin_mul_time;
    B_data[3] = sin_mul_time;
    B_data[4] = cos_mul_time;
    B_data[8] = time_s;

    /* 保存控制量 */
    for (int i = 0; i < 3; i++)
    {
        input_matrix_data[i] = input.data[i];
    }
    /* 计算状态预测量 */
    Kalman_Filter_Function::Forecast_State_Variable(sys_state_matrix, B_matrix, input_matrix, sys_state_forecast_matrix);

    /* 计算A矩阵 */
    A_data[2] = -input.v_y * cos_mul_time - input.v_x * sin_mul_time;
    A_data[5] = input.v_x * cos_mul_time - input.v_y * sin_mul_time;

    /* 计算W矩阵 */
    W_data[0] = cos_mul_time;
    W_data[1] = -sin_mul_time;
    W_data[3] = sin_mul_time;
    W_data[4] = cos_mul_time;
    W_data[8] = time_s;

    /* 计算转换后的执行噪声矩阵 */
    Kalman_Filter_Function::Cal_A_mul_B_A_Trans(W_matrix, R, R_matrix);

    /* 更新预测协方差矩阵 */
    Kalman_Filter_Function::Forecast_Covariance(A_matrix, covariance_matrix, R_matrix, covariance_forecast_matrix);

    if (no_measurement_flag)
    {
        /* 更新状态量和协方差 */
        for (int i = 0; i < 9; i++)
        {
            covariance_data[i] = covariance_forecast_data[i];
        }
        for (int i = 0; i < 3; i++)
        {
            sys_state_data.data[i] = sys_state_forecast_data.data[i];
        }
    }
    // LOG_D("w:%f,theta:%f no measure.", input.w, sys_state_data.theta);
}

/* 返回状态值 */
union Kalman_Filter_Coor_Namespace::Sys_State_Union Kalman_Filter_Coor_Namespace::Return_Coor(void)
{
    return sys_state_data;
}

void Kalman_Filter_Coor_Namespace::Set_Coor(const union Kalman_Filter_Coor_Namespace::Sys_State_Union &measure, bool is_dm)
{
    /* 直接修改坐标和协方差 */
    sys_state_data = measure;
    const arm_matrix_instance_f32 &Q = is_dm ? Q_DM_matrix : Q_TC_matrix;
    for (int i = 0; i < 9; i++)
    {
        covariance_data[i] = Q.pData[i];
    }
}
